import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from filterpy.monte_carlo import systematic_resample
from numpy.random import randn

# 从CSV加载真实期货数据
def generate_futures_data():
    """从data/SF0SH.csv加载真实期货价格数据"""
    df = pd.read_csv('data/SF0SH.csv')
    return pd.DataFrame({
        'date': pd.to_datetime(df['date_time']),
        'price': df['close']
    })

# 粒子滤波预测函数
def predict_futures_prices(data, n_particles=1000, prediction_days=30):
    """使用粒子滤波对期货价格进行预测"""
    # 初始化粒子
    initial_price = data['price'].iloc[0]
    particles = np.zeros((n_particles, 2))  # [价格, 价格变化率]
    particles[:, 0] = initial_price
    particles[:, 1] = np.random.normal(0, 0.01, n_particles)
    
    # 初始化权重
    weights = np.ones(n_particles) / n_particles
    
    # 观测噪声和过程噪声
    observation_noise = 0.1
    process_noise = 0.05
    
    # 历史数据拟合
    historical_predictions = []
    for price in data['price']:
        # 预测步骤 (状态转移)
        particles[:, 1] += randn(n_particles) * process_noise  # 变化率随机游走
        particles[:, 0] *= (1 + particles[:, 1])  # 价格更新
        
        # 更新步骤 (计算权重)
        error = price - particles[:, 0]
        weights *= np.exp(-0.5 * (error / observation_noise) ** 2)
        
        # 归一化权重
        weights += 1.e-300  # 避免数值下溢
        weights /= sum(weights)
        
        # 重采样 (当有效粒子数过低时)
        if 1. / np.sum(weights**2) < n_particles / 2:
            indexes = systematic_resample(weights)
            particles[:] = particles[indexes]
            weights.fill(1.0 / n_particles)
        
        # 记录当前预测
        mean_price = np.average(particles[:, 0], weights=weights)
        historical_predictions.append(mean_price)
    
    # 未来预测
    future_predictions = []
    last_particles = particles.copy()
    last_weights = weights.copy()
    
    for _ in range(prediction_days):
        # 预测下一天
        last_particles[:, 1] += randn(n_particles) * process_noise
        last_particles[:, 0] *= (1 + last_particles[:, 1])
        
        # 计算预测值
        mean_price = np.average(last_particles[:, 0], weights=last_weights)
        future_predictions.append(mean_price)
    
    # 构建完整的预测结果
    dates = pd.date_range(
        start=data['date'].iloc[0],
        periods=len(data) + prediction_days
    )
    all_predictions = historical_predictions + future_predictions
    
    return pd.DataFrame({
        'date': dates,
        'actual_price': list(data['price']) + [np.nan] * prediction_days,
        'predicted_price': all_predictions
    })

# 主函数
def main():
    # 加载真实数据
    data = generate_futures_data()
    
    # 进行预测
    predictions = predict_futures_prices(data, n_particles=1000, prediction_days=30)
    
    # 可视化结果
    plt.figure(figsize=(12, 6))
    plt.plot(predictions['date'], predictions['actual_price'], 'b-', label='real price')
    plt.plot(predictions['date'], predictions['predicted_price'], 'r--', label='prediction price')
    plt.axvline(x=data['date'].iloc[-1], color='g', linestyle='--', label='prediction start point')
    plt.title('期货价格预测')
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()
